﻿#include "bandsolver.h"
#include <fstream>
#include <iostream>
#include <algorithm>
#include <sys/stat.h>
#include "eigenhelper.h"
#include <sstream>

//获取文件大小
//由于MinGW 8 中std::filesystem的实现有BUG，所以这里用系统调用
size_t file_size(const std::string &path) {
    struct stat st{};
    auto ret = stat(path.c_str(), &st);
    return ret != 0 ? 0 : st.st_size;
}


constexpr bool BandSolver::datCheck(size_t fileSize, DatInfo info) {
    if (info.id != DAT_MAGIC) {
        return false;
    }
    auto aSize = fileSize - sizeof(DatInfo) - sizeof(float) * info.n;
    if (info.ver == DAT_COMPRESSED) {
        return aSize == SZ_FLOAT * (info.n) * (info.p + info.q + 1);
    } else if (info.ver == DAT_UNCOMPRESSED) {
        return aSize == SZ_FLOAT * (info.n) * (info.n);
    } else {
        return false;
    }
}


bool BandSolver::loadFromFile(const std::string &filePath) {
    DatInfo info{};
    std::ifstream datFile(filePath, std::ifstream::binary);
    if (!datFile.is_open()) {
        std::cout << "错误" << strerror(errno) << std::endl;
        return false;
    }
    auto fileSize = file_size(filePath);
    datFile.read(reinterpret_cast<char *>(&info), SZ_HEAD);
    if (!datCheck(fileSize, info)) {
        std::cout << "文件校验失败，请确认文件是否损坏！" << std::endl;
        return false;
    }
    m_n = info.n;
    m_p = info.p;
    m_q = info.q;
    Matf tempAf;
    Vecf tempBf(m_n);
    if (info.ver == DAT_UNCOMPRESSED) {
        tempAf.resize(m_n, m_n);
        datFile.read(reinterpret_cast<char *>(tempAf.data()), tempAf.size() * SZ_FLOAT);
        m_A = compress(m_p, m_q, tempAf.cast<double>());
    } else {
        tempAf.resize(m_n, m_p + m_q + 1);
        datFile.read(reinterpret_cast<char *>(tempAf.data()), tempAf.size() * SZ_FLOAT);
        m_A = tempAf.cast<double>();
    }
    datFile.read(reinterpret_cast<char *>(tempBf.data()), tempBf.size() * SZ_FLOAT);
    m_B = tempBf.cast<double>();
    return true;
}


Matd BandSolver::compress(size_t p, size_t q, const Matd &original) {
    size_t n = original.rows();
    Matd compressed(n, p + q + 1);
    Vecd temp(p + q + 1);
    for (size_t i = 0; i < p; i++) {
        temp.setZero().tail(q + i + 1) = original.row(i).head(q + i + 1);
        compressed.row(i) = temp;
    }
    for (size_t i = p; i < n - q; ++i) {
        compressed.row(i) = original.row(i).segment(i - p, p + q + 1).eval();
    }
    for (size_t i = n - q; i < n; ++i) {
        temp.setZero().head(n - i + p) = original.row(i).tail(n - i + p);
        compressed.row(i) = temp;
    }
    return compressed;
}


constexpr size_t BandSolver::lenA(const DatInfo info) {
    return (info.ver == DAT_COMPRESSED ? info.n * (info.p + info.q + 1) : info.n * info.n);
}


Vecd BandSolver::solve() {
    m_X = solveCompressed(m_p, m_q, m_A, m_B);
    return m_X;
}

void BandSolver::printAns(size_t num) {
    std::cout << ansToString(num) << std::endl;
}

std::string BandSolver::ansToString(size_t num) {
    std::stringstream s;
    num = num < m_X.size() ? num : m_X.size();
    s << m_X.head(num).transpose();
    return s.str();
}

Vecd BandSolver::solveCompressed(size_t p, size_t q, const Matd &a, const Vecd &b) {
    assert(a.rows() == b.size());
    assert(a.cols() == p + q + 1);
    size_t cols = a.cols();
    size_t n = a.rows();
    Matd A = a;
    Vecd B = b;
    Vecd temp(cols);
    double ratio;
    for (size_t j = 0; j < n; ++j) {
        //主元
        auto major = A(j, p);
        auto stop = std::min(j + p + 1, n);
        for (size_t k = j + 1; k < stop; ++k) {
            temp.setZero().segment(p + j - k, q + 1) = A.row(j).tail(q + 1);
            ratio = A(k, j + p - k) / major;
            A.row(k) -= ratio * temp;
            B[k] -= ratio * B[j];
        }
    }
    for (int i = static_cast<int>(n) - 1; i >= 0; --i) {
        B[i] /= A(i, p);
        for (int k = std::max(static_cast<int>(i - q), 0); k < i; ++k) {
            B[k] -= B[i] * A(k, i - k + p);
        }
    }
    return B;
}

